{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2827caf2",
   "metadata": {},
   "source": [
    "Transformer\n",
    "1、embedding部分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "3c5574b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2906, 0.9512, 0.0314, 0.2878],\n",
      "        [0.9940, 0.6749, 0.1298, 0.2266],\n",
      "        [0.8681, 0.1730, 0.1197, 0.2071],\n",
      "        [0.2888, 0.3239, 0.3211, 0.8634]])\n"
     ]
    }
   ],
   "source": [
    "#快捷键说明Ctrl+Alt+Enter允许代码段，Shift + Enter\t运行当前单元格，并跳转到下一个单元格\n",
    "import torch\n",
    "from torch import nn\n",
    "import math\n",
    "import torch.nn.functional as F\n",
    "#from d2l import torch as d2l\n",
    "random_torch=torch.rand(4,4)\n",
    "print(random_torch) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "591fbc09",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TokenEmbedding(nn.Embedding):\n",
    "    #词汇表大小，embedding的维度\n",
    "    def __init__(self,vocab_size,d_model):\n",
    "        super(TokenEmbedding,self).__init__(vocab_size.d_model,padding_idx=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96864520",
   "metadata": {},
   "source": [
    "![positionalembedding](./figure/positionalencoding.PNG)<br>\n",
    "i是一个词的某一个维度<br> d_model是一个词的维度<br> batch_size=32（32个句子），seq_len=10（每个句子10个词）x 是输入张量，通常是一个批量的词嵌入（Word Embedding）self.encoding[:seq_len, :] 从预计算的位置编码矩阵 self.encoding（形状 (max_len, d_model)）中截取前 seq_len 行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2204c2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "encoding= tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        ...,\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]], device='cuda:0') torch.Size([100, 512])\n",
      "pos= tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], device='cuda:0') torch.Size([100])\n",
      "pos= tensor([[0.],\n",
      "        [1.],\n",
      "        [2.],\n",
      "        [3.],\n",
      "        [4.],\n",
      "        [5.],\n",
      "        [6.],\n",
      "        [7.],\n",
      "        [8.],\n",
      "        [9.]], device='cuda:0') torch.Size([100, 1])\n",
      "_2i= tensor([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], device='cuda:0') torch.Size([256])\n",
      "self.encoding= tensor([[ 0.0000,  1.0000,  0.0000,  ...,  1.0000,  0.0000,  1.0000],\n",
      "        [ 0.8415,  0.5403,  0.8268,  ...,  1.0000,  0.0010,  1.0000],\n",
      "        [ 0.9093, -0.4161,  0.9302,  ...,  1.0000,  0.0021,  1.0000],\n",
      "        ...,\n",
      "        [ 0.3796, -0.9251,  0.1690,  ...,  0.9948,  0.0995,  0.9950],\n",
      "        [-0.5734, -0.8193,  0.9100,  ...,  0.9947,  0.1005,  0.9949],\n",
      "        [-0.9992,  0.0398,  0.8547,  ...,  0.9945,  0.1015,  0.9948]],\n",
      "       device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "class PositionEmbedding(nn.Module):\n",
    "    def __init__(self,d_model,max_len,device):\n",
    "        #直接调用父类函数的__init__()\n",
    "        #super(PositionEmbedding,self).__init__()\n",
    "        super().__init__()\n",
    "        self.encoding=torch.zeros(max_len,d_model,device=device)\n",
    "        print(\"encoding=\",self.encoding[:10],self.encoding.shape)\n",
    "        #不计算梯度\n",
    "        self.encoding.requires_grad=False\n",
    "        pos=torch.arange(0,max_len,device=device)\n",
    "        print(\"pos=\",pos[:10],pos.shape)\n",
    "        pos=pos.float().unsqueeze(dim=1)\n",
    "        print(\"pos=\",pos[:10],pos.  step=2,device=device).float()\n",
    "        print(\"_2i=\",_2i[:10],_2i.shape)\n",
    "        self.encoding[:,0::2]=torch.sin(pos/1000**(_2i/d_model))\n",
    "        self.encoding[:,1::2]=torch.cos(pos/1000**(_2i/d_model))\n",
    "        print(\"self.encoding=\",self.encoding)\n",
    "    def forward(self,x):\n",
    "        batch_size,seq_len=x.size()\n",
    "        return self.encoding[:seq_len,:]\n",
    "\n",
    "d_model=512\n",
    "max_len=100\n",
    "device=torch.device(\"cuda\"if torch.cuda.is_available() else \"cpu\")\n",
    "pos_embed=PositionEmbedding(d_model,max_len,device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8829b44d",
   "metadata": {},
   "source": [
    "1、单头单批量情况：\n",
    "输入句子：\"I love NLP\"\n",
    "\n",
    "seq_len = 3（句子中有 3 个 token）\n",
    "\n",
    "head_dim = 4（每个 token 被编码成 4 维向量）\n",
    "\n",
    "我们先只看 一个 head，一个 batch\n",
    "Q = [\n",
    "  [q11, q12, q13, q14],   ← \"I\"\n",
    "  [q21, q22, q23, q24],   ← \"love\"\n",
    "  [q31, q32, q33, q34]    ← \"NLP\"\n",
    "]\n",
    "Q：3*4\n",
    "K = [\n",
    "  [k11, k12, k13, k14],   ← \"I\"\n",
    "  [k21, k22, k23, k24],   ← \"love\"\n",
    "  [k31, k32, k33, k34]    ← \"NLP\"\n",
    "]\n",
    "K：3*4\n",
    "第一步：K.transpose(-2, -1)\n",
    "K^T = K.transpose(-2, -1) = [head_dim, seq_len] = [4, 3]\n",
    "\n",
    "K^T = [\n",
    "  [k11, k21, k31],\n",
    "  [k12, k22, k32],\n",
    "  [k13, k23, k33],\n",
    "  [k14, k24, k34]\n",
    "]\n",
    "K=4*3\n",
    "第二步：计算 Q @ K^T\n",
    "→ scores = Q @ K^T = [3, 3]\n",
    "scores = [\n",
    "  [s11, s12, s13],  ← \"I\" attends to: \"I\", \"love\", \"NLP\"\n",
    "  [s21, s22, s23],  ← \"love\" attends to: ...\n",
    "  [s31, s32, s33]   ← \"NLP\" attends to: ...\n",
    "]\n",
    "其中每个元素是一个点积：\n",
    "\n",
    "    s11 = Q[\"I\"] · K[\"I\"]\n",
    "\n",
    "    s12 = Q[\"I\"] · K[\"love\"]\n",
    "\n",
    "    s13 = Q[\"I\"] · K[\"NLP\"]\n",
    "\n",
    "    s21 = Q[\"love\"] · K[\"I\"]\n",
    "\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8999e0f3",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4bc69b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model,n_head):\n",
    "        super().__init__()\n",
    "        self.n_head=n_head\n",
    "        self.d_model=d_model\n",
    "        self.w_q=nn.Linear(d_model,d_model)\n",
    "        print(\"w_q=\",self.w_q)\n",
    "        self.w_k=nn.Linear(d_model,d_model)\n",
    "        self.w_v=nn.Linear(d_model,d_model)\n",
    "        self.w_combine=nn.Linear(d_model,d_model)\n",
    "        self.softmax=nn.Softmax(dim=1)\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "\n",
    "    def forward(self,q,k,v,mask=None):\n",
    "        batch_size,seq_len,dismesion=q.shape\n",
    "        # 1. 线性变换\n",
    "        #拆分的头数为512//2个头\n",
    "        head_dim=self.d_model//self.n_head\n",
    "        Q = self.w_q(q)  # (batch, seq_len, d_model)\n",
    "        K = self.w_k(k)  # (batch, seq_len, d_model)\n",
    "        V = self.w_v(v)  # (batch, seq_len, d_model)\n",
    "        # 2. 切分为多头\n",
    "        #表示：batch_size个批量，3 个词，每个词拆成 2 个头，每个头拥有256(head_dim)维的词向量。\n",
    "        #transpose(1, 2)将seq_len和n_head维度交换位置,变为2个词和3个头\n",
    "        #transpose 是为了把“每个词有多个头”的结构，变成“每个头看整句话”的结构\n",
    "        #从而让多个 attention head 并行计算，是高效实现多头注意力的标准做法。\n",
    "        Q = Q.view(batch_size, seq_len, self.n_head, head_dim).transpose(1, 2)  # (batch, n_head, seq_len, head_dim)\n",
    "        K = K.view(batch_size, seq_len, self.n_head, head_dim).transpose(1, 2)\n",
    "        V = V.view(batch_size, seq_len, self.n_head, head_dim).transpose(1, 2)\n",
    "        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(head_dim)\n",
    "        # 4. 掩码（可选）\n",
    "        if mask is not None:\n",
    "            scores = scores.masked_fill(mask == 0, float('-inf'))\n",
    "        # 5. Softmax和Dropout\n",
    "        attn_weights = self.softmax(scores)\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "        # 6. 加权求和\n",
    "        output = torch.matmul(attn_weights, V)  # (batch, n_head, seq_len, head_dim)\n",
    "        # 7. 合并多头\n",
    "        output = output.transpose(1, 2).contiguous()  # (batch, seq_len, n_head, head_dim)\n",
    "        output = output.view(batch_size, seq_len, self.d_model)  # (batch, seq_len, d_model)\n",
    "        # 8. 最终线性变换\n",
    "        output = self.w_combine(output)\n",
    "        \n",
    "        return output\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c6bfa36",
   "metadata": {},
   "source": [
    "2、多头注意力机制"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7a1f670",
   "metadata": {},
   "source": [
    "![lr](./figure/lineartrans.PNG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "617e1041",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "w_q= Linear(in_features=512, out_features=512, bias=True)\n"
     ]
    }
   ],
   "source": [
    "x=torch.rand(128,32,512)\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model,n_head):\n",
    "        super().__init__()\n",
    "        self.n_head=n_head\n",
    "        self.d_model=d_model\n",
    "        self.w_q=nn.Linear(d_model,d_model)\n",
    "        print(\"w_q=\",self.w_q)\n",
    "        self.w_k=nn.Linear(d_model,d_model)\n",
    "        self.w_v=nn.Linear(d_model,d_model)\n",
    "        self.w_combine=nn.Linear(d_model,d_model)\n",
    "        self.softmax=nn.Softmax(dim=1)\n",
    "\n",
    "    def forward(self,q,k,v,mask=None):\n",
    "        batch,time,dismesion=q.shape\n",
    "        n_d=d_model//self.n_head\n",
    "        q,k,v=self.w_q(q),self.w_k(k),self.w_v(v)\n",
    "        q=q.view(batch,time,self)\n",
    "        score=q@k.transpose(2,3)/math.sqrt(n_d)\n",
    "        \n",
    "  \n",
    "\n",
    "d_model = 512\n",
    "n_head = 8\n",
    "mha = MultiHeadAttention(d_model, n_head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4743407b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "2d01bb77",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.2812, -2.1981, -0.6320,  0.9519],\n",
      "         [ 0.1983,  1.5821, -1.6014,  0.4388],\n",
      "         [-0.3685, -0.5047,  0.2372, -1.3599]],\n",
      "\n",
      "        [[ 0.6744,  1.8030, -1.2131,  0.8643],\n",
      "         [-0.0467, -0.6300,  0.0264,  1.1371],\n",
      "         [ 0.3850,  1.7408,  1.0076, -0.7746]]]) torch.Size([2, 3, 4])\n",
      "tensor([[[ 0.2812,  0.1983, -0.3685],\n",
      "         [-2.1981,  1.5821, -0.5047],\n",
      "         [-0.6320, -1.6014,  0.2372],\n",
      "         [ 0.9519,  0.4388, -1.3599]],\n",
      "\n",
      "        [[ 0.6744, -0.0467,  0.3850],\n",
      "         [ 1.8030, -0.6300,  1.7408],\n",
      "         [-1.2131,  0.0264,  1.0076],\n",
      "         [ 0.8643,  1.1371, -0.7746]]]) torch.Size([2, 4, 3])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn(2, 3, 4)  # 形状：(batch_size=2, seq_len=3, features=4)\n",
    "print(x,x.shape)\n",
    "y = x.permute(0, 2, 1)  # 新形状：(2, 4, 3)\n",
    "print(y,y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "8913362c",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'builtin_function_or_method' object is not subscriptable",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTypeError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[46]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m F.one_hot(\u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mtensor\u001b[49m\u001b[43m[\u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[32;43m2\u001b[39;49m\u001b[43m]\u001b[49m,\u001b[32m1\u001b[39m) \n",
      "\u001b[31mTypeError\u001b[39m: 'builtin_function_or_method' object is not subscriptable"
     ]
    }
   ],
   "source": [
    "F.one_hot(torch.tensor[0,2],1) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b17dae4",
   "metadata": {},
   "source": [
    "LayerNorm层实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "db090aac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入数据形状: torch.Size([2, 3, 4])\n",
      "gamma: Parameter containing:\n",
      "tensor([1., 1., 1., 1.], requires_grad=True)\n",
      "beta: Parameter containing:\n",
      "tensor([0., 0., 0., 0.], requires_grad=True)\n",
      "输出形状: torch.Size([2, 3, 4])\n",
      "输出数据:\n",
      " tensor([[[-1.1619, -0.3873,  0.3873,  1.1619],\n",
      "         [-1.1619, -0.3873,  0.3873,  1.1619],\n",
      "         [-1.1619, -0.3873,  0.3873,  1.1619]],\n",
      "\n",
      "        [[ 1.1619,  0.3873, -0.3873, -1.1619],\n",
      "         [ 1.1619,  0.3873, -0.3873, -1.1619],\n",
      "         [ 1.1619,  0.3873, -0.3873, -1.1619]]], grad_fn=<AddBackward0>)\n",
      "更新后的 gamma: Parameter containing:\n",
      "tensor([0.9904, 1.0000, 1.0000, 1.0096], requires_grad=True)\n",
      "更新后的 beta: Parameter containing:\n",
      "tensor([-0.0100, -0.0100, -0.0100, -0.0100], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn,optim\n",
    "x = torch.tensor([\n",
    "    [[1.0, 2.0, 3.0, 4.0],   # 样本1的第1个时间步\n",
    "     [0.0, 1.0, 2.0, 3.0],   # 样本1的第2个时间步\n",
    "     [-1.0, 0.0, 1.0, 2.0]], # 样本1的第3个时间步\n",
    "     \n",
    "    [[4.0, 3.0, 2.0, 1.0],   # 样本2的第1个时间步\n",
    "     [3.0, 2.0, 1.0, 0.0],   # 样本2的第2个时间步\n",
    "     [2.0, 1.0, 0.0, -1.0]]  # 样本2的第3个时间步\n",
    "], dtype=torch.float32)\n",
    "\n",
    "print(\"输入数据形状:\", x.shape) \n",
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self,d_model, eps=1e-6):#数值稳定性\n",
    "        super().__init__()\n",
    "        self.gamma=nn.Parameter(torch.ones(d_model))  # 可学习参数 nn.Parameter 是 PyTorch 中的一个核心类，\n",
    "        #​用于将张量（Tensor）标记为模型的可学习参数。它的主要作用是告诉 PyTorch 的优化器（如 Adam、SGD）在训练过程中需要更新这些张量的值（即通过梯度下降进行优化）。\n",
    "        self.beta=nn.Parameter(torch.zeros(d_model))\n",
    "        self.eps = eps\n",
    "    def forward(self, x):\n",
    "        mean=x.mean(dim=-1, keepdim=True)  # 计算均值维度不变\n",
    "        var=x.var(dim=-1,keepdim=True)  # 计算方差维度不变\n",
    "        out=(x-mean)/torch.sqrt(var+self.eps)  # 标准化\n",
    "        out=self.gamma*out+self.beta\n",
    "        return out\n",
    "    \n",
    "\n",
    "d_model = 4  # 特征维度需与输入数据的最后一维一致\n",
    "layernorm = LayerNorm(d_model=d_model, eps=1e-6)\n",
    "optimizer = optim.Adam(layernorm.parameters(), lr=0.01) \n",
    "\n",
    "# 打印可学习参数 gamma 和 beta\n",
    "print(\"gamma:\", layernorm.gamma)  # 初始值为全1\n",
    "print(\"beta:\", layernorm.beta)  \n",
    "output1 = layernorm(x)\n",
    "print(\"输出形状:\", output.shape)  # 输出: torch.Size([2, 3, 4])\n",
    "print(\"输出数据:\\n\", output)\n",
    "loss = output1.sum()\n",
    "optimizer.zero_grad()      # 清空梯度\n",
    "loss.backward()\n",
    "optimizer.step()\n",
    "\n",
    "output2 = layernorm(x)\n",
    "print(\"更新后的 gamma:\", layernorm.gamma) \n",
    "print(\"更新后的 beta:\", layernorm.beta) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
